import pickle
import sys
import time
from tqdm import tqdm
import numpy as np
import os
import torch
import torch.optim as optim

class Experiment():

    def __init__(self, option, learner, data):

        self.option = option
        self.learner = learner
        self.data = data

        self.learning_rate = option.learning_rate
        self.grad_clip = option.max_grad_norm

        # helpers
        self.msg_with_time = lambda msg: \
            "%s Time elapsed %0.2f hrs (%0.1f mins)" \
            % (msg, (time.time() - self.start) / 3600.,
               (time.time() - self.start) / 60.)

        self.start = time.time()
        self.epoch = 0
        self.best_valid_loss = np.inf
        self.best_valid_in_top = 0.
        self.train_stats = []
        self.valid_stats = []
        self.test_stats = []
        self.early_stopped = False
        self.log_file = open(os.path.join(self.option.this_expsdir, "log.txt"), "w")

        self.optimizer = optim.Adam(self.learner.parameters(), lr=self.learning_rate)

    def clip_gradient(self, optimizer, grad_clip):
        """
        Clips gradients computed during backpropagation to avoid explosion of gradients.

        :param optimizer: optimizer with the gradients to be clipped
        :param grad_clip: clip value
        """
        for group in optimizer.param_groups:
            for param in group["params"]:
                if param.grad is not None:
                    param.grad.data.clamp_(-grad_clip, grad_clip)

    def one_epoch(self, mode, num_batch, next_fn):
        epoch_loss = []
        epoch_in_top = []
        for batch in range(num_batch):

            if (batch + 1) % max(1, (int(num_batch / self.option.print_per_batch))) == 0:
                sys.stdout.write("%d/%d\t" % (batch + 1, num_batch))
                sys.stdout.flush()

            (qq, hh, tt), mdb = next_fn()
            if len(qq) != self.option.batch_size*2:
                print('Validation - Input Shape Issue:', len(qq))
                continue

            if mode == "train":
                batch_loss, final_loss, in_top, predictions = self.learner(qq, hh, tt, mdb)
                batch_loss.backward()
                # torch.nn.utils.clip_grad_norm_(self.learner.parameters(), self.grad_clip)
                self.clip_gradient(self.optimizer, self.grad_clip)
                self.optimizer.step()
                self.learner.zero_grad()
            else:
                with torch.no_grad():
                    batch_loss, final_loss, in_top, predictions = self.learner(qq, hh, tt, mdb)

            epoch_loss += list(final_loss.cpu().data.tolist())
            epoch_in_top += list(in_top.cpu().data.tolist())

        msg = self.msg_with_time(
            "Epoch %d mode %s Loss %0.4f In top %0.4f."
            % (self.epoch + 1, mode, np.mean(epoch_loss), np.mean(epoch_in_top)))
        print(msg)
        self.log_file.write(msg + "\n")

        return epoch_loss, epoch_in_top

    def one_epoch_train(self):
        if self.epoch > 0 and self.option.resplit:
            self.data.train_resplit(self.option.no_link_percent)
        loss, in_top = self.one_epoch("train",
                                      self.data.num_batch_train,
                                      self.data.next_train)

        self.train_stats.append([loss, in_top])

    def one_epoch_valid(self):
        loss, in_top = self.one_epoch("valid",
                                      self.data.num_batch_valid,
                                      self.data.next_valid)
        self.valid_stats.append([loss, in_top])
        self.best_valid_loss = min(self.best_valid_loss, np.mean(loss))
        self.best_valid_in_top = max(self.best_valid_in_top, np.mean(in_top))

    def one_epoch_test(self):
        loss, in_top = self.one_epoch("test",
                                      self.data.num_batch_test,
                                      self.data.next_test)
        self.test_stats.append([loss, in_top])

    def early_stop(self):
        loss_improve = self.best_valid_loss == np.mean(self.valid_stats[-1][0])
        in_top_improve = self.best_valid_in_top == np.mean(self.valid_stats[-1][1])
        if loss_improve or in_top_improve:
            return False
        else:
            if self.epoch < self.option.min_epoch:
                return False
            else:
                return True

    def train(self):
        while (self.epoch < self.option.max_epoch and not self.early_stopped):
            self.one_epoch_train()
            self.one_epoch_valid()
            self.one_epoch_test()
            self.epoch += 1
            model_path = '{}-{}.pth'.format(self.option.model_path, self.epoch)
            torch.save(self.learner.state_dict(), model_path)
            print("Model saved at %s" % model_path)

            if self.early_stop():
                self.early_stopped = True
                print("Early stopped at epoch %d" % (self.epoch))

        all_test_in_top = [np.mean(x[1]) for x in self.test_stats]
        best_test_epoch = np.argmax(all_test_in_top)
        best_test = all_test_in_top[best_test_epoch]

        msg = "Best test in top: %0.4f at epoch %d." % (best_test, best_test_epoch + 1)
        print(msg)
        self.log_file.write(msg + "\n")
        pickle.dump([self.train_stats, self.valid_stats, self.test_stats],
                    open(os.path.join(self.option.this_expsdir, "results.pkl"), "wb"))

    def get_predictions(self):
        if self.option.query_is_language:
            all_accu = []
            all_num_preds = []
            all_num_preds_no_mistake = []

        f = open(os.path.join(self.option.this_expsdir, "test_predictions.txt"), "w")
        if self.option.get_phead:
            f_p = open(os.path.join(self.option.this_expsdir, "test_preds_and_probs.txt"), "w")
        all_in_top = []
        for batch in range(self.data.num_batch_test):
            if (batch + 1) % max(1, (self.data.num_batch_test / self.option.print_per_batch)) == 0:
                sys.stdout.write("%d/%d\t" % (batch + 1, self.data.num_batch_test))
                sys.stdout.flush()
            (qq, hh, tt), mdb = self.data.next_test()

            with torch.no_grad():
                _, _, in_top, predictions_this_batch = self.learner(qq, hh, tt, mdb)  # (b, num_entity)
            all_in_top += list(in_top)

            for i, (q, h, t) in enumerate(zip(qq, hh, tt)):
                p_head = predictions_this_batch[i, h]
                if self.option.adv_rank:
                    eval_fn = lambda j, p: p >= p_head and (j != h)
                    this_predictions = filter(eval_fn, enumerate(predictions_this_batch[i, :]))
                elif self.option.rand_break:
                    eval_fn = lambda j, p: (p > p_head) or (
                                (p == p_head) and (j != h) and (np.random.uniform() < 0.5))
                    this_predictions = filter(eval_fn, enumerate(predictions_this_batch[i, :]))
                else:
                    eval_fn = lambda p: p > p_head
                    tmp = predictions_this_batch[i, :].unbind()
                    this_predictions = filter(eval_fn, tmp)

                this_predictions = sorted(this_predictions, key=lambda x: x[1], reverse=True)
                if self.option.query_is_language:
                    all_num_preds.append(len(this_predictions))
                    mistake = False
                    for k, _ in this_predictions:
                        assert (k != h)
                        if not self.data.is_true(q, k, t):
                            mistake = True
                            break
                    all_accu.append(not mistake)
                    if not mistake:
                        all_num_preds_no_mistake.append(len(this_predictions))
                else:
                    this_predictions.append((h, p_head))
                    this_predictions = [self.data.number_to_entity[j] for j, _ in this_predictions]
                    q_string = self.data.parser["query"][q]
                    h_string = self.data.number_to_entity[h]
                    t_string = self.data.number_to_entity[t]
                    to_write = [q_string, h_string, t_string] + this_predictions
                    f.write(",".join(to_write) + "\n")
                    if self.option.get_phead:
                        f_p.write(",".join(to_write + [str(p_head)]) + "\n")
        f.close()
        if self.option.get_phead:
            f_p.close()

        if self.option.query_is_language:
            print("Averaged num of preds", np.mean(all_num_preds))
            print("Averaged num of preds for no mistake", np.mean(all_num_preds_no_mistake))
            msg = "Accuracy %0.4f" % np.mean(all_accu)
            print(msg)
            self.log_file.write(msg + "\n")

        msg = "Test in top %0.4f" % np.mean(all_in_top)
        msg += self.msg_with_time("\nTest predictions written.")
        print(msg)
        self.log_file.write(msg + "\n")